from ultralytics import YOLO




if __name__ == '__main__':

    # Load a model
    model = YOLO("yolov5s_under_water.yaml")  #  这边加载模型的yaml
    model = YOLO("yolov5su.pt")  # 这边加载预训练权重文件
    model = YOLO("yolov5s_under_water.yaml").load("yolov5su.pt")  # build from YAML and transfer weights

    # Train the model
    results = model.train(data="under-water.yaml", epochs=10, imgsz=640,device=0)    #data 部分写数据集配置文件   即under-water.yaml epoch。batchsize都可根据需要调，开启gpu训练 设置device=0